import os
from datetime import datetime

import pandas as pd
from dateutil.relativedelta import relativedelta

from app_config import get_engine_ts
from app_config import get_pro


def offset_date(original_date_str):
    # 将输入的日期字符串转换为日期对象
    original_date = datetime.strptime(original_date_str, '%Y%m%d')

    # 计算6个月后的日期
    new_date = original_date + relativedelta(months=6)

    # 将结果日期格式化为 'yyyyMMdd' 字符串
    return new_date.strftime('%Y%m%d')


def calc_kc50(strStartDate='', strEndDate='', str_filePre='', str_000688Date=''):

    if not os.path.exists(str_filePre):
        os.makedirs(str_filePre)
        print(f"文件夹 '{str_filePre}' 创建成功")
    engine = get_engine_ts()

    """     
            （1）上市时间超过一个季度，除非该证券自上市以来的日均总市值在沪深市场中排名前 30 位；
            （2）非 ST、*ST 证券    
    """
    kcb = get_pro().stock_basic(market='科创板', list_status='L', fields='ts_code,symbol,name,area,industry,list_date')
    cyb = get_pro().stock_basic(market='创业板', list_status='L', fields='ts_code,symbol,name,area,industry,list_date')
    kcb_cyb = pd.concat([kcb, cyb], ignore_index=True)
    # kcOneYear = kcb_cyb[kcb_cyb['list_date'] < offset_date(strStartDate)]
    ts_codes = kcb_cyb['ts_code'].tolist()

    # 构建查询语句
    query = f"""
    SELECT ts_code,trade_date, amount FROM `daily`
    WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
    AND trade_date >= '{strStartDate}'
    AND trade_date <= '{strEndDate}'
    """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['amount'] = result_df['amount'] / 1000
    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', '日均成交额']
    merge = pd.merge(kcb_cyb, grouped_df, on='ts_code', how='left')

    # 构建查询语句
    query = f"""
        SELECT ts_code,trade_date, total_mv FROM `daily_basic`
        WHERE ts_code IN ({','.join(f"'{code}'" for code in ts_codes)})
        AND trade_date >= '{strStartDate}'
        AND trade_date <= '{strEndDate}'
        """
    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['total_mv'] = result_df['total_mv'] / 10000
    grouped_df_total_mv = result_df.groupby('ts_code')['total_mv'].mean().reset_index()
    grouped_df_total_mv.columns = ['ts_code', '日均总市值']
    pd_merge = pd.merge(merge, grouped_df_total_mv, on='ts_code', how='left')

    """
        过去一年日均成交金额排名位于样本空间前 80%。
    """
    df_sorted = pd_merge.sort_values(by='日均成交额', ascending=False).reset_index(drop=True)
    df_sorted['i_avgAmount'] = df_sorted.index
    # 获取前 90% 的数据
    num_rows = len(df_sorted)
    top_80_percent = df_sorted.head(int(num_rows * 0.8))

    """
        将待选样本按照过去一年日均总市值由高到低排名，选取排名前 50 的证券作为指数样本
    """
    top_80_percent = top_80_percent.sort_values('日均总市值', ascending=False).reset_index(drop=True)
    top_80_percent['i_totalMv'] = top_80_percent.index
    df_result_kc50 = top_80_percent.head(120)

   
    df_930955 = get_pro().index_weight(index_code='931643.CSI', start_date=str_000688Date, end_date=str_000688Date)
    df_930955.rename(columns={'con_code': 'ts_code'}, inplace=True)
    df_930955.rename(columns={'trade_date': 'list_date'}, inplace=True)
    df_930955['symbol'] = df_930955['ts_code'].str.split('.').str[0]
    df_930955.drop(columns=['weight'], inplace=True)
    df_930955.to_excel(str_filePre + str_000688Date + '-931643.xlsx')

    # 获取 dataframe1 中的所有列
    columns_dataframe1 = df_result_kc50.columns

    # 创建缺失列并填充为 '*'
    for column in columns_dataframe1:
        if column not in df_930955.columns:
            df_930955[column] = '**'

    # 将 dataframe2 的列顺序调整为与 dataframe1 一致
    df_930955 = df_930955[columns_dataframe1]

    # 追加 dataframe2 到 dataframe1
    result = pd.concat([df_result_kc50, df_930955], ignore_index=True)

    result.to_excel(str_filePre + '科创创业50_2024_.xlsx')


if __name__ == '__main__':
    """2024年三季度"""
    strStartDate = '20230801'
    strEndDate = '20240731'
    str_filePre = 'zfile/科创创业50_2024季度3/'
    str_000688Date = '20240628'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # """2024年2季度"""
    strStartDate = '20230501'
    strEndDate = '20240430'
    str_filePre = 'zfile/科创创业50_2024季度2/'
    str_000688Date = '20240628'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # """2024年1季度"""
    strStartDate = '20230201'
    strEndDate = '20240131'
    str_filePre = 'zfile/科创创业50_2024季度1/'
    str_000688Date = '20240430'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # """2023年四季度"""
    strStartDate = '20221101'
    strEndDate = '20231031'
    str_filePre = 'zfile/科创创业50_2023季度4/'
    str_000688Date = '20240131'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # """2023年三季度"""
    strStartDate = '2022801'
    strEndDate = '20230731'
    str_filePre = 'zfile/科创创业50_2023季度3/'
    str_000688Date = '20231031'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # """2023年二季度"""
    strStartDate = '2022501'
    strEndDate = '20230430'
    str_filePre = 'zfile/科创创业50_2023季度2/'
    str_000688Date = '20230731'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    ##
    #  上: 上市时间调整为超过 12 个月
    #  下: 上市时间超过 6 个月
    ##

    # """2023年一季度"""
    strStartDate = '20220201'
    strEndDate = '20230131'
    str_filePre = 'zfile/科创创业50_2023季度1/'
    str_000688Date = '20230428'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # """2022年四季度"""
    strStartDate = '20211101'
    strEndDate = '20221031'
    str_filePre = 'zfile/科创创业50_2022季度4/'
    str_000688Date = '20230131'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

    # """2022年三季度"""
    strStartDate = '20210801'
    strEndDate = '20220731'
    str_filePre = 'zfile/科创创业50_2022季度3/'
    str_000688Date = '20221031'
    calc_kc50(strStartDate=strStartDate, strEndDate=strEndDate, str_filePre=str_filePre, str_000688Date=str_000688Date)

